package com.ktjiaoyu.travelaiagent.advisors;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.MessageAggregator;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * 顾问类，用于在请求AI大模型时进行Token使用统计
 * 参考SimpleLoggerAdvisor实现自定义顾问逻辑。
 * @see org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor
 */
@Slf4j
public class TokenCounterAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {
    // 用于统计各个会话使用的Token总数量
    private static final Map<String, Long>  TOKEN_USAGE_MAP = new ConcurrentHashMap() ;

    private AdvisedRequest before(AdvisedRequest request) {
        // 在请求AI大模型之前，可以在这里进行一些前置处理，比如记录日志、鉴权等
        return request;
    }
    private void observeAfter(AdvisedResponse advisedResponse) {
        // 从Advisor上下文中可以获取到会话ID
        // 前提是设置了参数：spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
        String chatId = (String)advisedResponse.adviseContext()
                .get(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY);
        // 本次请求使用的token总数
        Long tokenUsed = advisedResponse.response().getMetadata().getUsage().getTotalTokens().longValue();
        // 计算总Token数并保存
        Long totalTokens = countTokens(chatId, tokenUsed);
        // 输出Token当前使用量和总使用量
        log.info("ChatId: {}, Current token used: {} , Total token used: {}", chatId, tokenUsed, totalTokens);
    }

    /** 统计某个会话的总token用量 */
    private static Long countTokens(String chatId, Long tokenUsed) {
        // 先根据会话ID获得总token使用量
        Long totalTokens = TOKEN_USAGE_MAP.get(chatId);
        if(totalTokens == null){
            // 如果是第一次对话，则总数就是当前对话的token数量
            totalTokens = tokenUsed;
        }else{
            // 如果不是第一次对话，则追加token数量
            totalTokens += tokenUsed;
        }
        // 将计算后的总数添加到Map中
        TOKEN_USAGE_MAP.put(chatId, totalTokens);
        return totalTokens;
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        // 类似与AOP中的环绕增强，先执行before
        advisedRequest = before(advisedRequest);
        // 责任链模式，调用下一个顾问，或者请求AI大模型获得结果
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
        // 执行后置任务处理响应
        observeAfter(advisedResponse);
        // 最后返回响应
        return advisedResponse;
    }
    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        advisedRequest = this.before(advisedRequest);
        Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);
        return (new MessageAggregator()).aggregateAdvisedResponse(advisedResponses, this::observeAfter);
    }
    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }
    @Override
    public int getOrder() {
        return 100;
    }
}
